import numpy as np
import open3d as o3d

from cuboid_locate import CuboidSurface


def convert_xyz_to_pcd(xyz_file_path: str):

    coordinates = list()
    with open(xyz_file_path) as xyz_file:
        for line in xyz_file:
            x, y, z = [float(e) for e in line.rstrip().split(",")]
            coordinates.append([x, y, z])
    coordinates = np.asarray(coordinates)

    pcd = o3d.geometry.PointCloud()
    pcd.points = o3d.utility.Vector3dVector(coordinates)

    save_path = xyz_file_path.replace(".xyz", ".pcd")
    o3d.io.write_point_cloud(save_path, pcd)


def draw_cuboid_corners(pcd: o3d.geometry.PointCloud, surface: CuboidSurface):
    pcd.paint_uniform_color([0.5, 0.5, 0.5])

    top_left = surface.top_left
    top_right = surface.top_right
    bottom_left = surface.bottom_left
    bottom_right = surface.bottom_right

    pcd_tree = o3d.geometry.KDTreeFlann(pcd)

    top_left = np.array([top_left.x, top_left.y, top_left.z])
    [k, idx, _] = pcd_tree.search_knn_vector_3d(top_left, 20)
    np.asarray(pcd.colors)[idx[1:], :] = [1, 0, 0]
    top_right = np.array([top_right.x, top_right.y, top_right.z])
    [k, idx, _] = pcd_tree.search_knn_vector_3d(top_right, 20)
    np.asarray(pcd.colors)[idx[1:], :] = [0, 1, 0]
    bottom_left = np.array([bottom_left.x, bottom_left.y, bottom_left.z])
    [k, idx, _] = pcd_tree.search_knn_vector_3d(bottom_left, 20)
    np.asarray(pcd.colors)[idx[1:], :] = [0, 0, 1]
    bottom_right = np.array([bottom_right.x, bottom_right.y, bottom_right.z])
    [k, idx, _] = pcd_tree.search_knn_vector_3d(bottom_right, 20)
    np.asarray(pcd.colors)[idx[1:], :] = [0, 1, 1]
